from colorama import Fore
import random
import sys
from aalpy.base.SUL import SUL
from aalpy.base import Oracle
from learning.DBhelper import *
from learning.Errors import NonDeterministicError
import utils
from logger import GLOBAL_LOGGER as LOG


###
# The code used in this file is copied from the AALpy project:
# https://github.com/DES-Lab/AALpy
#
# Following file/class has been copied:
# -- aalpy/oracles/StatePrefixEqOracle.py
#
# Adaptions to the existing code have been made:
# -- check for non-determinism
# -- check for connection errors
#
#
###

MAX_NONDETERMINISTIC_ERROR_TIMES = 20

class StatePrefixOracleFailSafe(Oracle):

    def __init__(self, alphabet: list, sul: SUL, walks_per_state=10, walk_len=30, depth_first=False, database=''):
        super().__init__(alphabet, sul)
        self.walks_per_state = walks_per_state
        self.steps_per_walk = walk_len
        self.depth_first = depth_first
        self.freq_dict = dict()
        self.dbhelper = DBhelper(database)
        # 人工输入的反例列表：存储待测试的输入序列
        self.manual_counterexamples = []

    def add_manual_counterexample(self, input_sequence):
        """
        添加人工输入的反例（输入序列）
        
        参数:
            input_sequence: 输入序列，可以是列表或元组，例如 ['ClientHello', 'Certificate', 'Finish']
        """
        if isinstance(input_sequence, (list, tuple)):
            # 转换为元组格式
            sequence_tuple = tuple(input_sequence) if isinstance(input_sequence, list) else input_sequence
            self.manual_counterexamples.append(sequence_tuple)
            LOG.logger.info(f'Added manual counterexample: {sequence_tuple}')
        else:
            LOG.logger.warning(f'Invalid input sequence format: {input_sequence}')
    
    def test_sequence_from_s0(self, input_sequence):
        """
        从 s0（初始状态）开始测试输入序列，返回输出序列
        
        参数:
            input_sequence: 输入序列，可以是列表或元组
            
        返回:
            output_sequence: 输出序列
        """
        if isinstance(input_sequence, list):
            input_sequence = tuple(input_sequence)
        
        LOG.logger.info(f'Testing sequence from s0: {input_sequence}')
        
        # 重置 SUL 到初始状态
        self.sul.pre()
        
        output_sequence = []
        try:
            for input_symbol in input_sequence:
                output = self.sul.step(input_symbol)
                output_sequence.append(output)
            
            LOG.logger.info(f'Output sequence: {output_sequence}')
            return tuple(output_sequence)
        except Exception as e:
            LOG.logger.error(f'Error testing sequence: {e}')
            return None
        finally:
            # 清理
            self.sul.post()
    
    def find_cex(self, hypothesis):
        LOG.logger.info('start to find counterexample') 
        LOG.logger.info('='*100)
        
        # 首先检查人工输入的反例
        if self.manual_counterexamples:
            LOG.logger.info('Checking manual counterexamples...')
            for manual_cex in self.manual_counterexamples:
                LOG.logger.info(f'Testing manual counterexample: {manual_cex}')
                
                # 从 s0 开始测试这个序列
                self.reset_hyp_and_sul(hypothesis)
                output_sequence = []
                hypothesis_outputs = []
                
                try:
                    # 测试 SUL 的输出
                    for input_symbol in manual_cex:
                        sul_output = self.sul.step(input_symbol)
                        output_sequence.append(sul_output)
                        
                        # 测试 hypothesis 的输出
                        hyp_output = hypothesis.step(input_symbol)
                        hypothesis_outputs.append(hyp_output)
                        
                        # 如果发现不一致，返回这个反例
                        if sul_output != hyp_output:
                            LOG.logger.info(f'Found counterexample from manual input: {manual_cex}')
                            LOG.logger.info(f'SUL outputs: {output_sequence}')
                            LOG.logger.info(f'Hypothesis outputs: {hypothesis_outputs}')
                            # 保存到数据库
                            self.dbhelper.insert_query(str(manual_cex), str(output_sequence))
                            return manual_cex
                    
                    # 如果完全一致，记录日志但继续
                    LOG.logger.info(f'Manual counterexample did not find inconsistency: {manual_cex}')
                    LOG.logger.info(f'SUL outputs: {output_sequence}')
                    LOG.logger.info(f'Hypothesis outputs: {hypothesis_outputs}')
                    # 仍然保存到数据库
                    self.dbhelper.insert_query(str(manual_cex), str(output_sequence))
                except Exception as e:
                    LOG.logger.warning(f'Error testing manual counterexample {manual_cex}: {e}')
                    continue
            
            # 清空已处理的人工反例（可选：如果希望重复使用，可以注释掉这行）
            # self.manual_counterexamples = []
        
        states_to_cover = []
        for state in hypothesis.states:
            if state.prefix not in self.freq_dict.keys():
                self.freq_dict[state.prefix] = 0

            states_to_cover.extend([state] * (self.walks_per_state - self.freq_dict[state.prefix]))

        if self.depth_first:
            # reverse sort the states by length of their access sequences
            # first do the random walk on the state with longest access sequence
            states_to_cover.sort(key=lambda x: len(x.prefix), reverse=True)
        else:
            random.shuffle(states_to_cover)
        
        for state in states_to_cover:
            LOG.logger.info('*'*100)
            LOG.logger.info('current state : ' + str(state.state_id))
            self.freq_dict[state.prefix] += 1

            non_det_attempts = 0
            
            suffix_buf = ()
            for _ in range(self.steps_per_walk):
                suffix_buf += (random.choice(self.alphabet),)
            
            LOG.logger.info('prefix : ' + str(state.prefix))
            LOG.logger.info('suffix : ' + str(suffix_buf))
            results_in_db_str = self.dbhelper.execute_query(str(state.prefix + suffix_buf))
            if results_in_db_str:
                LOG.logger.info('found in sqlite : ' + results_in_db_str)
                results_in_db = str2list(results_in_db_str)
                hypothesis.reset_to_initial()
                query = state.prefix + suffix_buf
                out_hyps = []
                out_suls = []
                current_query = ()
                for i in range(len(query)):
                    current_query += query[i:i+1]
                    out_hyp = hypothesis.step(query[i])
                    out_hyps.append(out_hyp)
                    out_sul = results_in_db[i]
                    out_suls.append(out_sul)
                    if out_sul != out_hyp:
                        LOG.logger.info('found counterexample : ' + str(current_query))
                        LOG.logger.info('out_hyps : ' + str(out_hyps))
                        LOG.logger.info('out_suls : ' + str(out_suls))
                        return current_query
                continue
            
            NonDeterministicError_times = 0
            while True:
                try:
                    is_prefix_consistent = True
                    self.reset_hyp_and_sul(hypothesis)
                    out_sul = ''
                    out_suls = []
                    out_hyp = ''
                    out_hyps = []
                    prefix = state.prefix
                    for p in prefix:
                        self.num_steps += 1
                        out_hyp = hypothesis.step(p)
                        out_hyps.append(out_hyp)
                        out_sul = self.sul.step(p)
                        out_suls.append(out_sul)
                        if out_sul != out_hyp:
                            non_det_attempts += 1
                            is_prefix_consistent = False
                            break
                    if not is_prefix_consistent:
                        continue
                        
                    suffix = ()
                    for i in range(len(suffix_buf)):
                        suffix += suffix_buf[i:i+1]
                        self.num_steps += 1
                        out_sul = self.sul.step(suffix[-1])
                        out_suls.append(out_sul)
                        out_hyp = hypothesis.step(suffix[-1])
                        out_hyps.append(out_hyp)
                        
                        if out_sul != out_hyp:
                            LOG.logger.info('found counterexample : ' + str(prefix + suffix))
                            LOG.logger.info('out_hyps : ' + str(out_hyps))
                            LOG.logger.info('out_suls : ' + str(out_suls))
                            self.dbhelper.insert_query(str(prefix + suffix), str(out_suls))
                            LOG.logger.info('insert to database')
                            return prefix + suffix
                        
                    self.dbhelper.insert_query(str(prefix + suffix), str(out_suls)) 
                    LOG.logger.info('insert to database')
                    break
                except NonDeterministicError as e:
                    NonDeterministicError_times += 1
                    if NonDeterministicError_times >= MAX_NONDETERMINISTIC_ERROR_TIMES:
                        LOG.logger.error(Fore.RED + 'Reach max nondeterministic error times.')
                        sys.exit(-1)
                    LOG.logger.warning(e.message)    
                    continue

        return None